#############################################
# Object detection - YOLO - OpenCV
# Author : Arun Ponnusamy   (July 16, 2018)
# Website : http://www.arunponnusamy.com
############################################
import os

import cv2
import numpy as np
from PIL import Image

import tools.infer.__utility__.utility_barcode as utility
from tools.infer.__ocr__.barcode_ocr import TextDetector, predict_img
from tools.infer.predict_rec import TextRecognizer
ori_path = 'D:/program/opencv_yolo/object-detection-opencv'



def get_output_layers(net):
    
    layer_names = net.getLayerNames()
    
    output_layers = [layer_names[i[0] - 1] for i in net.getUnconnectedOutLayers()]

    return output_layers


def draw_prediction(img, class_id, confidence, x, y, x_plus_w, y_plus_h):

    label = str(classes[class_id])

    color = COLORS[class_id]

    cv2.rectangle(img, (x,y), (x_plus_w,y_plus_h), color, 2)
    txt = label +" "+ str(round(confidence, 3))
    cv2.putText(img, txt, (x-10,y-10),cv2.FONT_HERSHEY_SIMPLEX, 0.5, color, 2)
    #cv的图像操作


scale = 0.00392

classes = None

with open(os.path.join(ori_path,"cfg/barcode.names"), 'r') as f:
    classes = [line.strip() for line in f.readlines()]

COLORS = np.random.uniform(0, 255, size=(len(classes), 3))




##################################################################################################
#运行需要 权值、设置、类



net = cv2.dnn.readNet(os.path.join(ori_path,"backup/barcode_last.weights"),os.path.join(ori_path, "cfg/barcode.cfg"))



def test_and_show(path, savepath):
    print("##################################")
    args = utility.parse_args()
    print(args)



    text_detector = TextDetector(args)
    text_recognizer = TextRecognizer(args)


    img_file_path = 'D:/program/xmltotxt/xsz/Images'
    img_list = os.listdir(img_file_path)




    for i in img_list:
        name = i
        image = cv2.imread(os.path.join(img_file_path,i))

        img = Image.open(os.path.join(img_file_path,i))
        img.show()
        Width = image.shape[1]
        Height = image.shape[0]
        blob = cv2.dnn.blobFromImage(image, scale, (416, 416), (0, 0, 0), True, crop=False)

        net.setInput(blob)

        outs = net.forward(get_output_layers(net))

        class_ids = []
        confidences = []
        boxes = []
        conf_threshold = 0.5
        nms_threshold = 0.4

        for out in outs:
            for detection in out:
                scores = detection[5:]
                class_id = np.argmax(scores)
                confidence = scores[class_id]
                if confidence > 0.5:
                    center_x = int(detection[0] * Width)
                    center_y = int(detection[1] * Height)
                    w = int(detection[2] * Width)
                    h = int(detection[3] * Height)
                    x = center_x - w / 2
                    y = center_y - h / 2
                    class_ids.append(class_id)
                    confidences.append(float(confidence))
                    boxes.append([x, y, w, h])
                    cv2.rectangle(image, (x, y), (x + w, y + h), (0, 255, 0), 1)
                    cv2.putText(image, name, (x - 10, y - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 0), 2)
                    cv2.imshow("src",image)
        #print(boxes)
        indices = cv2.dnn.NMSBoxes(boxes, confidences, conf_threshold, nms_threshold)
        #print("indices", indices)

        for j in indices:
            print("j", j )
            i = j[0]
            box = boxes[i]
            x = int(box[0])
            y = int(box[1])
            w = int(box[2])
            h = int(box[3])
            '''
            if class_ids[i] == 1:
                imgcut = image[y:y + h, x:x + w]  ## 裁剪坐标为[y0:y1, x0:x1]
                # if name == "handwriting":
                #  cv2.imwrite(".", imgcut)
                # cv2.imshow("result", imgcut)
                # cv2.waitKey()
                # 二值化 adaptiveThreshold()
                imgcut = cv2.cvtColor(imgcut, cv2.COLOR_BGR2GRAY)
                thrpic = cv2.adaptiveThreshold(imgcut, 255, cv2.ADAPTIVE_THRESH_GAUSSIAN_C,
                                               cv2.THRESH_BINARY, 31, 11)
                cv2.imwrite("./handwriting/"+name, thrpic)
                '''
            img = image[y:y+h, x:x+w]
            cv2.imshow("cut",img)
            rec = predict_img(img, text_detector,text_recognizer)

            print("results:",rec)

test_and_show("./xsz", "./error")
